Description

Adds TPU support for context_parallel_strategy=ring in Flash attention. This PR uses and adapts some Tokamax Splash components and includes the MaxText integration, config validation, and sharding checks.

Key Changes:

• TPU ring context parallelism: context_parallel_strategy=ring now runs on TPU (previously GPU only). Q remains sequence sharded while KV rotate around the context mesh axis, partial results are combined with online softmax.
• Adapted Tokamax Splash components: the experimental TPU Splash attention family is adapted into src/maxtext/kernels/tokamax_splash_attention/, using Tokamax 0.0.12 (openxla/tokamax@4936e75). This replaces the previous imports from tokamax._src.ops.experimental... and provides the necessary modifications the ring path needs.
• Integration layer: a new tokamax_ring_attention.py module owns the MaxText/Tokamax interface such as mesh-axis derivation, sharding checks (ring axis only on the sequence dim, GQA/MQA head mapping preserved, head/dKV dims local), SplashConfig construction, and the sharded ring kernel call.
• Config validation: the TPU ring is enforced at config time so that unsupported combinations can fail early with explicit errors instead of at trace or compile time.
• Tests: ring forward/gradient vs a single-device Splash reference, mask for the adapted components, config-validation cases, and an e2e attention test under a CP mesh.

Key files:

• src/maxtext/kernels/attention/tokamax_ring_attention.py: This detects when ring CP is requested, derives Q/KV mesh axes, validates ring-mesh/head/dKV sharding, builds the Tokamax SplashConfig and lazy causal mask, and wraps the kernel with make_sharded_ring_attention_kernel / call_ring_attention.
• src/maxtext/kernels/tokamax_splash_attention/ (adapted from Tokamax):
  • splash_attention_kernel.py: the Splash kernel, adapted to provide log-sum-exp and max logits, converted from base-2 to natural base, and to share types via base.py.
  • ring_attention_kernel.py: the ring forward/backward over KV shards adapted to factor reusable logic into ring_attention_utils.py and to validate the ring-axis size against the built kernel.
  • ring_attention_utils.py: shared ring helpers such as MaskInfo slicing, Q-sequence offsets per KV shard, block-activity checks.
  • splash_attention_mask.py / splash_attention_mask_info.py: mask classes and block-metadata construction, adapted for slicing mask metadata per ring step.
  • base.py: shared types (SegmentIds, residuals, mask constants).
• src/maxtext/layers/attention_op.py: adds mesh/sharding validation when ring is requested, rejects features the ring path does not support like sinks, indexer masks, ragged attention, chunked prefill, bidirectional masks, max-logit recording...; routes the flash path to the sharded ring kernel when ring is requested.
• src/maxtext/configs/types.py: accepts context_parallel_strategy=ring on TPU and enforces the supported combination likeattention=flash + use_tokamax_splash=True, rejects packing, CP load balancing, ragged attention, sinks, indexer, chunked prefill, MoBA, multimodal, QK-clip, and dropout with explicit errors.
• docs/guides/optimization/sharding.md: ring CP is now supported on TPU via the Tokamax Splash path (in addition to the GPU Transformer Engine path).
• tests/unit/tokamax_ring_attention_test.py: ring output and gradients vs Splash reference.
• tests/unit/tokamax_splash_attention_mask_test.py: tests for the adapted mask components.
• tests/unit/attention_test.py: e2eAttentionOp ring-CP test under a CP mesh.
• tests/unit/configs_value_test.py: validation cases for the TPU ring.

Limitation and Future work

• Support load balancing and sequence packing

Performance comparison

• As noted, more detailed performance comparison will be provided when load balancing is supported
• Compile time vs baseline all-gather

Run command

Example on a v5p-16 (8 chips) and 64K context:

python3 -m maxtext.trainers.pre_train.train \
src/maxtext/configs/base.yml \
base_output_directory=gs://<GCS_BUCKET> run_name=ring_cp_demo \
model_name=llama3.1-8b steps=10 enable_checkpointing=false dataset_type=synthetic \
max_target_length=65536 per_device_batch_size=0.125 \
ici_context_parallelism=8 \
attention=flash use_tokamax_splash=True \
context_parallel_strategy=ring \
context_parallel_load_balance=False packing=False

Tests

Ran:

python3 -m pytest tests/unit/tokamax_ring_attention_test.py
python3 -m pytest tests/unit/tokamax_splash_attention_mask_test.py
python3 -m pytest tests/unit/configs_value_test.py -k tpu_tokamax_ring
python3 -m pytest tests/unit/attention_test.py -k tpu_flash_attention_ring_context_parallel

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.